// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: junxiao, leon

#pragma once

#include <algorithm>
#include <memory>
#include <set>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include <queue>

#include <glog/logging.h>
#include "coin/proto/coin_order_enums.pb.h"
#include "coin2/app/driver.h"
#include "coin2/base/log.h"
#include "coin2/exchange/base/api_util/auth_key.h"
#include "coin2/exchange/base/feed/update.h"
#include "coin2/strategy/dex_strat/simple_arb_strat_config.pb.h"
#include "coin2/strategy/order_executor/agg_executor.h"
#include "coin2/strategy/strategy.h"
#include "coin2/strategy/util/moving_ohlc.h"
#include "coin2/strategy/util/latency_util.h"

namespace coin2::strategy::dex_strat {

using coin::proto::OrderSide;
using coin2::strategy::AggressiveExecutorConfig;
using coin2::strategy::PassiveExecutorConfig;
using ::coin2::strategy::order_executor::AggOrderExecutor;
using ::coin2::app::impl::Driver;
using ::coin2::exchange::base::feed::FeedUpdate;
using ::coin2::exchange::base::market::MarketExchangeApi;
using ::coin2::exchange::base::order::IOrderSubsystem;
using ::coin2::exchange::base::order::OrderUpdate;
using ::coin2::exchange::base::symbology::IProductHolder;
using ::coin::proto::OrderEvent;
using ::coin2::strategy::util::MovingHigh;
using ::coin2::strategy::util::MovingLow;
using ::coin2::strategy::dex_strat::RefProductConvertDirection;

class SimpleArbSubStrategy;

struct SimpleArbTickEntry {
  int64_t trade_ts;
  double trade_ask0;
  double trade_bid0;
  int64_t ref_ts;
  double ref_ask0;
  double ref_bid0; 
  SimpleArbTickEntry(int64_t trade_ts = 0, double trade_ask0 = 0.0, double trade_bid0 = 0.0, int64_t ref_ts = 0, double ref_ask0 = 0.0, double ref_bid0 = 0.0)
    : trade_ts(trade_ts), trade_ask0(trade_ask0), trade_bid0(trade_bid0), ref_ts(ref_ts), ref_ask0(ref_ask0), ref_bid0(ref_bid0) {}
};

class SimpleArbPricer {
 public:
  explicit SimpleArbPricer(
    const IProduct* trade_product,
    const IProduct* main_ref_product,
    RefProductConvertDirection ref_convert_direction,
    const IProduct* convert_ref_product,
    int64_t max_queue_size = 1,
    double arb_diff_bar = 2.5,
    double arb_leave_bar = -0.5,
    double arb_stop_price_diff = 3.0,
    int64_t tick_ts_diff_allowed = 300'000'000LL
    );
  void onBookFeed(const std::string& product_key, const FeedUpdate& upd);
  void onPosChanged(int pos_side, double order_price);
  std::pair<bool, std::pair<double, OrderSide>> CalculatePrice();
  bool Ready();
 private:
  std::function<const IProductHolder*(int64_t)> get_product_holder_;
  std::string trade_product_key_;
  std::string main_ref_key_;
  std::string convert_ref_key_;
  RefProductConvertDirection ref_product_convert_direction_{RefProductConvertDirection::NONE};
  std::optional<int64_t> last_main_ref_ts_;
  std::optional<PriceQty> last_main_ref_bid0_;
  std::optional<PriceQty> last_main_ref_ask0_;
  std::optional<int64_t> last_convert_ref_ts_;
  std::optional<PriceQty> last_convert_ref_bid0_;
  std::optional<PriceQty> last_convert_ref_ask0_;
  const int64_t max_queue_size_;
  const double arb_diff_bar_;
  const double arb_leave_bar_;
  const double arb_stop_price_diff_;
  const int64_t tick_ts_diff_allowed_;
  std::deque<SimpleArbTickEntry> signal_queue_;
  int current_pos_side_;
  double order_price_;
  std::pair<bool, std::pair<double, OrderSide>> current_op_;

  bool IsStartSellEntry(const SimpleArbTickEntry& tick_entry);
  bool IsStartBuyEntry(const SimpleArbTickEntry& tick_entry);
  bool IsReturnSellEntry(const SimpleArbTickEntry& tick_entry);
  bool IsReturnBuyEntry(const SimpleArbTickEntry& tick_entry);
  bool ShouldStartSell();
  bool ShouldStartBuy();
  bool ShouldReturnSell();
  bool ShouldReturnBuy();
  void CalculateCurrentOperation();
};

class SimpleArbStrategy : public ::coin2::strategy::IStrategy {
 public:
  SimpleArbStrategy(
      const std::experimental::filesystem::path& strat_json,
      bool dry_run,
      std::optional<int64_t> timestamp);
  void Init(::coin2::app::Driver* driver) override;
  void InitAuthKey(::coin2::app::Driver* driver);
  void InitReserveMap();
  void onBookFeed(const FeedUpdate& upd) override;
  void onAccountInfo(const OrderUpdate& upd) override;
  void onAccountOrder(const OrderUpdate& upd) override;
  void onOrderLog(const coin::proto::OrderGatewayLog& order_log) override;
  void InjectFeed(const coin2::feed::mgr::impl::FeedSubsystemManager* feed_mgr) final;
 private:
  int64_t strategy_created_time_;
  SimpleArbProductConfig config_;
  std::vector<std::unique_ptr<SimpleArbSubStrategy>> simple_arb_sub_strats_;
  std::unordered_map<std::string, std::vector<SimpleArbSubStrategy*>> feed_sub_map_;
  std::unordered_map<std::string, SimpleArbSubStrategy*> order_sub_map_;
  std::unique_ptr<StrategyReporter2> strat_reporter_;
  std::unique_ptr<AuthKey> auth_key_;
};

class SimpleArbSubStrategy : public ::coin2::strategy::IStrategy {
 public:
  using ::coin2::strategy::IStrategy::Init;
  using ::coin2::strategy::IStrategy::onBookFeed;
  SimpleArbSubStrategy(
      SimpleArbStrategy* parent,
      const SimpleArbProductConfig& config,
      int64_t strategy_created_time,
      bool enable_telemetry,
      bool dry_run);
  MarketExchangeApi mea() { return MarketExchangeApi::FromString(config_.trade_mea()); }
  IOrderSubsystem* os() { return driver_->order()->GetUniqueSystem(mea()); }
  void Init(::coin2::app::Driver* driver,
            ::coin2::strategy::util::StrategyReporter2* reporter);
  const SimpleArbProductConfig& GetConfig() const { return config_; }
  const IProduct* GetTradeProduct() const {return trade_product_.get(); }
  const IProduct* GetTradeProductForReporter() const {return trade_product_for_reporter_.get(); }
  const IProduct* GetMainRefProduct() const {return main_ref_product_.get(); }
  const IProduct* GetConvertRefProduct() const {return convert_ref_product_.get(); }
  AggOrderExecutor* GetMutableExecutor() { return executor_.get(); }
  void onAccountInfo(const OrderUpdate& upd) override;
  void onBookFeed(const std::string& product_key, const FeedUpdate& upd);
  void InjectFeed(const coin2::feed::mgr::impl::FeedSubsystemManager* feed_mgr) final;

 protected:
  SimpleArbStrategy* parent_{nullptr};
  SimpleArbProductConfig config_;
  int64_t strategy_created_time_;
  std::unique_ptr<IProduct> trade_product_;
  std::unique_ptr<IProduct> trade_product_for_reporter_;
  std::unique_ptr<IProduct> main_ref_product_;
  std::unique_ptr<IProduct> convert_ref_product_;
  RefProductConvertDirection ref_product_convert_direction_{RefProductConvertDirection::NONE};
  std::unique_ptr<AuthKey> auth_key_;
  Driver* driver_;
  exchange::base::strategy_util::StrategyLogger* strat_logger_;
  std::unique_ptr<AggOrderExecutor> executor_;
  std::unique_ptr<SimpleArbPricer> pricer_;
  bool enable_telemetry_{false};
  bool dry_run_{false};
  std::set<std::string> feed_keys_;
  std::set<std::string> ref_feed_keys_;
  const coin2::feed::mgr::impl::FeedSubsystemManager* feed_mgr_{nullptr};

  double reserved_pos_;
  double lot_size_;
  double max_pos_;
  double min_pos_;

  double pos_mismatch_bar_;
  int64_t order_cooldown_;
  int64_t safe_order_cooldown_;
  double order_tolerated_ratio_;
  int64_t tick_ts_diff_allowed_;
  double arb_diff_bar_;
  double arb_leave_bar_;
  double arb_stop_price_diff_;
  int64_t signal_queue_size_;

  int64_t trade_tick_cnt_;
  double current_pos_;
  bool pos_initialized_;
  bool is_pos_updating_;
  bool is_last_order_return_type_;
  int64_t last_order_ts_;
  double last_order_price_;
  void UpdatePosition();
};

}  // namespace coin2::strategy::dex_strat
